from iotdb_api import query
from iotdb_api import get_timestamp
import json


year = 2024

def query_size_month(year,month_begin, month_end):
    """
    查询并返回每个月的出境数据流量总和
    :param month_begin: 开始月份
    :param month_end: 结束月份
    :return: 该月出境数据流量总和
    """
    if month_begin != "12":
        sql_str = f"select data_size from root.all_msg where time > {year}-{month_begin}-01T00:00:00.000 and time < {year}-{month_end}-01T00:00:00.000"
    else:
        sql_str = f"select data_size from root.all_msg where time > {year}-{month_begin}-01T00:00:00.000 and time < {year+1}-{month_end}-01T00:00:00.000"
    data = query(sql_str)
    is_empty = len(data['values'])

    if is_empty != 0:
        sum_size = 0
        data_len = len(data['values'][0])
        for i in range(data_len):
            sum_size += float(data['values'][0][i][:-3])
        return sum_size

    pass

def query_year(year):
    """
    以月为单位，提供一年的出境数据
    :param year: 年份
    :return: json 数据
    """
    # 数据定义
    #month_lst = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']  # 月份英文缩写
    lst_size = []
    #data_dic = {'月份': month_lst, }
    # 获取每个月的数据流量总和
    for i in range(1, 13):
        if i != 12:
            month_begin = str(i).rjust(2, "0")
            month_end = str(i + 1).rjust(2, "0")
        else:
            month_begin = "12"
            month_end = "01"
        size = query_size_month(year,month_begin, month_end)
        lst_size.append(size)
        #data_dic.update({'出境流量': lst_size})

    # 返回 json 格式数据
    #data_json = json.dumps(data_dic, indent=4)
    data_dic =  {}
    data1 = []
    data_dic["year"] = year
    for i in range(12):
        data2 = {}
        data2["month"] = i + 1
        data2["value"] = lst_size[i]
        data1.append(data2)
    data_dic["data"] = data1
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_type(month,type):
    """
    查询每个处境模式下的数据流量出境情况
    :param month: 月份
    :param type: 出境模式
    :return:
    """
    if month !=12:
        sql_str = f"select data_size from root.all_msg where model = '{type}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select data_size from root.all_msg where model = '{type}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
    data = query(sql_str)
    is_empty = len(data['values'])

    if is_empty != 0:
        sum_size = 0
        data_len = len(data['values'][0])
        for i in range(data_len):
            sum_size += float(data['values'][0][i][:-3])
        return sum_size

    pass


def query_type(month):
    """
    以出境模式为单位，提供一个月的出境数据流量
    :param month: 月份
    :return: json 数据
    """
    data_dic = {} 
    type_lst = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
    data_dic["month"] = month
    data1 = []
    for i in range(len(type_lst)):
        data2 = {}
        size = query_size_type(month,type_lst[i])
        data2["type"] = type_lst[i]
        data2["value"] = size
        data1.append(data2)

    data_dic["data"] = data1
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_level(month,level):
    """
    查询每个敏感等级下的数据流量出境情况
    :param month: 月份
    :param level: 敏感等级
    :return:
    """
    if month !=12:
        sql_str = f"select data_size from root.all_msg where level = '{level}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select data_size from root.all_msg where level = '{level}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
    data = query(sql_str)
    is_empty = len(data['values'])

    if is_empty != 0:
        sum_size = 0
        data_len = len(data['values'][0])
        for i in range(data_len):
            sum_size += float(data['values'][0][i][:-3])
        return sum_size

    pass

def query_level(month):
    """
    以敏感等级为单位，提供一个月的出境数据流量
    :param month: 月份
    :return: json 数据
    """
    data_dic = {}  # K-V : 敏感等级-流量总和
    level_lst = [0,1,2,3]
    data_dic["month"] = month
    data1 = []
    # 获取每个敏感等级的数据流量总和
    for i in range(len(level_lst)):
        data2 = {}
        size = query_size_level(month,level_lst[i])
        data2["level"] = level_lst[i]
        data2["value"] = size
        #data_dic.update({level_lst[i]: size})
        data1.append(data2)

    data_dic["data"] = data1
    # 返回 json 格式数据
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_nation(month):
    """
    查询 一个月内的各个国家数据流量情况
    :param month: 月份
    :return:
    """
    if month !=12:
        sql_str = f"select nation from root.all_msg where nation is not null and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select nation from root.all_msg where nation is not null and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
    data_prov = query(sql_str)
    is_empty = len(data_prov['values'])

    if is_empty != 0:
        nation_lst = list(set(data_prov['values'][0]))
        data_len = len(nation_lst)
        data_dic = {}
        data_dic["month"] = month
        data1 = []
        for i in range(data_len):
            data2 = {}
            data2["nation"] = nation_lst[i]
            if month!=12:
                sql_str = f"select data_size from root.all_msg where nation = '{nation_lst[i]}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
            else:
                sql_str = f"select data_size from root.all_msg where nation = '{nation_lst[i]}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
            data = query(sql_str)
            sum_size = 0
            length = len(data['values'][0])
            for j in range(length):
                sum_size += float(data['values'][0][j][:-3])
            data2["value"] = sum_size
            data1.append(data2)
        data_dic["data"]=data1
        return data_dic

    pass


def query_nation(month):
    data_dic = query_size_nation(month)
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_province(month):
    """
    查询 一个月内的各个省份数据出境情况
    :param month: 月份
    :return:
    """
    if month !=12:
        sql_str = f"select prov_id,dst_ip from root.province where time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select prov_id,dst_ip from root.province where time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
    data_prov = query(sql_str)
    is_empty = len(data_prov['values'])

    if is_empty != 0:
        prov_lst = list(set(data_prov['values'][0]))
        dip_lst = list(set(data_prov['values'][1]))
        data_len = len(prov_lst)
        #data_dic = {
        #    '行政省市': prov_lst,
        #    '主要出境国家': dip_lst}
        data_dic = {}
        data_dic["month"] = month
        data1 = []
        for i in range(data_len):
            data2 = {}
            data2["name"] = prov_lst[i]
            if month!=12:
                sql_str = f"select data_size from root.province where prov_id = '{prov_lst[i]}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
            else:
                sql_str = f"select data_size from root.province where prov_id = '{prov_lst[i]}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
            data = query(sql_str)
            lst_size = []
            lst_ip = []
            sum_size = 0
            length = len(data['values'][0])
            for j in range(length):
                sum_size += float(data['values'][0][j][:-3])
            #lst_size.append(sum_size)
            data2["value"] = sum_size
            data1.append(data2)
            #data_dic.update({"出境流量": lst_size})
        data_dic["data"]=data1
        return data_dic

    pass


def query_province(month):
    data_dic = query_size_province(month)
    data_json = json.dumps(data_dic)
    return data_json

    pass




def query_all(month):
    data_dic = {}
    data_dic["month"] = month
    data_dic["province"] = query_size_province(month)["data"]
    level_lst = [0,1,2,3]
    data1 = []
    for i in range(len(level_lst)):
        data2 = {}
        size = query_size_level(month,level_lst[i])
        data2["level"] = level_lst[i]
        data2["value"] = size
        data1.append(data2)
    data_dic["level"] = data1
    type_lst = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
    data1 = []
    for i in range(len(type_lst)):
        data2 = {}
        size = query_size_type(month,type_lst[i])
        data2["type"] = type_lst[i]
        data2["value"] = size
        data1.append(data2)
    data_dic["type"] = data1
    data_dic["nation"] = query_size_nation(month)["data"]
    data_json = json.dumps(data_dic)
    return data_json

companys = [
    "aaa","bbb","ccc","ddd","eee","fff","ggg","hhh",
    "iii","jjj","kkk","lll","mmm","nnn","ooo","ppp",
]

def query_size_company(month):
    """
    查询一个月企业数据出境情况
    :param month: 月份
    :return:
    """
    if month!=12:
        sql_str = f"select src_id from root.company where time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select src_id from root.company where time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"

    data_id = query(sql_str)
    is_empty = len(data_id['values'])

    if is_empty != 0:
        company_lst = list(set(data_id['values'][0]))
        data_len = len(company_lst)
        data_dic = {}
        data_dic["month"] = month
        data1 = []
        for i in range(data_len):
            data2 = {}
            data2["id"] = company_lst[i]
            data2["name"] = "公司"+str(data2["id"])
            if month!=12:
                sql_str = f"select data_size,level from root.company where src_id = '{company_lst[i]}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
            else:
                sql_str = f"select data_size,level from root.company where src_id = '{company_lst[i]}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
            data_query = query(sql_str)
            length = len(data_query['values'][0])
            data_size = data_query['values'][0]
            level = data_query['values'][1]
            size = 0
            size_level_zero = 0
            size_level_one = 0
            size_level_two = 0
            size_level_three = 0
            for j in range(length):
                if level[j] == "0":
                    size_level_zero += float(data_size[j][:-3])
                elif level[j] == "1":
                    size_level_one += float(data_size[j][:-3])
                elif level[j] == "2":
                    size_level_two += float(data_size[j][:-3])
                elif level[j] == "3":
                    size_level_three += float(data_size[j][:-3])
                size += float(data_size[j][:-3])
            data2["value"] = size
            data3 = []
            data3.append({"level":0,"value":size_level_zero})
            data3.append({"level":1,"value":size_level_one})
            data3.append({"level":2,"value":size_level_two})
            data3.append({"level":3,"value":size_level_three})
            data2["data"] = data3
            data1.append(data2)
        data1 = sorted(data1, key=lambda item: item['value'], reverse = True)
        data1 = data1[:20]
        data_dic["data"] = data1
        return data_dic

    pass


def query_company(month):
    data_dic = query_size_company(month)
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_company_one(id,month):
    """
    查询一个月企业数据出境情况
    :param id: 企业号
    :param month: 月份
    :return:
    """

    if month!=12:
        sql_str = f"select data_size,level,model,subsidiary from root.company where src_id = '{id}' and time > {year}-{month}-01T00:00:00.000 and time < {year}-{month+1}-01T00:00:00.000"
    else:
        sql_str = f"select data_size,level,model,subsidiary from root.company where src_id = '{id}' and time > {year}-12-01T00:00:00.000 and time < {year+1}-01-01T00:00:00.000"
    data_query = query(sql_str)
    try:
        length = len(data_query['values'][0])
        data_size = data_query['values'][0]
        level = data_query['values'][1]
        type = data_query['values'][2]
        subsidiary = data_query['values'][3]
    except:
        length = 0
    data_dic = {}
    data_dic["id"] = id
    data_dic["month"] = month
    data_dic["value"] = 0
    subsidiary_list = [
        {"name":"子公司1","value":0},
        {"name":"子公司2","value":0},
        {"name":"子公司3","value":0},
        {"name":"子公司4","value":0},
        {"name":"子公司5","value":0}
    ]
    level_list = [
        {"name":0,"value":0},
        {"name":1,"value":0},
        {"name":2,"value":0},
        {"name":3,"value":0}
    ]
    type_list = [
        {"name":1,"value":0},
        {"name":2,"value":0},
        {"name":3,"value":0},
        {"name":4,"value":0},
        {"name":5,"value":0},
        {"name":6,"value":0},
        {"name":7,"value":0},
        {"name":8,"value":0},
        {"name":9,"value":0},
        {"name":10,"value":0},
        {"name":11,"value":0},
        {"name":12,"value":0},
        {"name":13,"value":0},
        {"name":14,"value":0},
        {"name":15,"value":0}
    ]

    for j in range(length):
        subsidiary_list[int(subsidiary[j])-1]["value"] += float(data_size[j][:-3])
        level_list[int(level[j])]["value"] += float(data_size[j][:-3])
        type_list[int(type[j])-1]["value"] += float(data_size[j][:-3])
        data_dic["value"] += float(data_size[j][:-3])

    data_dic["Subsidiary"] = subsidiary_list
    data_dic["level"] = level_list
    data_dic["type"] = type_list
    return data_dic

    pass


def query_company_one(id,month):
    data_dic = query_size_company_one(id,month)
    data_json = json.dumps(data_dic)
    return data_json

    pass

def query_size_route(id):
    data_dic ={}
    data_dic["id"] = id
    sql_str = f"select route from root.all_msg where iden_id = '{id}'"
    data_query = query(sql_str)
    try:
        routes = data_query['values'][0][0].split('.')
        data_dic["source"] = routes[0]
        data_dic["Destination"] = routes[len(routes)-1]
        data1 = []
        for i in range(len(routes)):
            data2 = {}
            data2["id"] = i + 1
            data2["name"] = routes[i]
            data1.append(data2)
        data_dic["process"] = data1
    except:
        data_dic["source"] = ""
        data_dic["Destination"] = ""
        data_dic["process"] = []
    return data_dic

def query_route(id):
    data_dic = query_size_route(id)
    data_json = json.dumps(data_dic)
    return data_json

    pass

# if __name__ == '__main__':
#     query_size() # 1.1 接口
#     query_level() # 1.3 接口
#     query_company() # 2 接口
#     query_province() # 3 接口
